/******************************************************************************
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * libkperf licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Author:
 * Create: 2025-01-13
 * Description: Thread binding system for performance optimization
 ******************************************************************************/
#ifndef THREAD_BINDING_H
#define THREAD_BINDING_H

#include <vector>
#include <string>
#include <unordered_map>
#include <memory>
#include <chrono>
#include <atomic>
#include <thread>
#include <mutex>
#include <condition_variable>

// Forward declarations
struct PmuData;
struct PmuAttr;

namespace ThreadBinding {

/**
 * @brief Thread performance metrics structure
 */
struct ThreadMetrics {
    int pid;                    // Process ID
    unsigned tid;               // Thread ID
    unsigned llc_miss;          // LLC miss count
    unsigned llc_cache;         // LLC cache access count
    double miss_ratio;          // LLC miss ratio
    double activity_score;      // Activity score (0.0-1.0)
    std::string comm;           // Process name
    std::vector<int> cpu_affinity; // Original CPU affinity list
    int numa_id;                // NUMA node ID
    std::chrono::steady_clock::time_point last_update; // Last update timestamp

    ThreadMetrics(int p, unsigned t, unsigned miss, unsigned cache, 
                  const std::string& process_name = "", 
                  const std::vector<int>& affinity = {})
        : pid(p), tid(t), llc_miss(miss), llc_cache(cache), 
          comm(process_name), cpu_affinity(affinity), numa_id(-1),
          last_update(std::chrono::steady_clock::now()) {
        miss_ratio = (cache != 0) ? static_cast<double>(miss) / cache : 0.0;
        activity_score = calculateActivityScore();
    }

    bool operator<(const ThreadMetrics& other) const {
        return activity_score > other.activity_score;
    }

    double calculateActivityScore() const;
};

/**
 * @brief CPU core information structure
 */
struct CpuCore {
    int core_id;                // Core ID
    int numa_id;                // NUMA node ID
    bool is_available;          // Whether core is available for binding
    std::chrono::steady_clock::time_point last_binding; // Last binding time

    CpuCore(int id, int numa) : core_id(id), numa_id(numa), is_available(true),
        last_binding(std::chrono::steady_clock::now()) {}
};

/**
 * @brief NUMA node information structure
 */
struct NumaNode {
    int numa_id;                // NUMA node ID
    std::vector<int> cpu_cores; // CPU cores in this NUMA
    std::vector<int> available_cores; // Available cores for binding

    NumaNode(int id) : numa_id(id) {}
};

/**
 * @brief Binding strategy enumeration
 */
enum class BindingStrategy {
    ACTIVITY_BASED,     // Bind based on activity score
    MISS_RATIO_BASED,   // Bind based on LLC miss ratio
    HYBRID              // Hybrid approach
};

/**
 * @brief Configuration structure
 */
struct Config {
    std::string process_keyword;    // Process keyword to match
    uint64_t collection_interval_ms; // Data collection interval (ms)
    uint64_t binding_interval_ms;   // Binding update interval (ms)
    uint64_t min_collection_time_ms; // Minimum collection time (ms)
    BindingStrategy strategy;       // Binding strategy
    bool enable_continuous_binding; // Enable continuous binding updates
    int max_threads_per_core;       // Maximum threads per core
    bool verbose;                   // Verbose output

    Config() : collection_interval_ms(100), binding_interval_ms(1000),
               min_collection_time_ms(100), strategy(BindingStrategy::ACTIVITY_BASED),
               enable_continuous_binding(true), max_threads_per_core(1), verbose(false) {}
};

/**
 * @brief Thread binding manager class
 */
class ThreadBindingManager {
public:
    explicit ThreadBindingManager(const Config& config);
    ~ThreadBindingManager();

    // Main interface
    bool initialize();
    void start();
    void stop();
    bool is_running() const { return running_.load(); }

    // Configuration
    void update_config(const Config& config);
    const Config& get_config() const { return config_; }

    // Statistics
    struct Statistics {
        uint64_t total_bindings;
        uint64_t successful_bindings;
        uint64_t failed_bindings;
        uint64_t collection_cycles;
        std::chrono::steady_clock::time_point start_time;
        std::chrono::steady_clock::time_point last_update;
    };

    Statistics get_statistics() const;

    // Thread binding operations
    bool perform_strategy_specific_binding(std::vector<ThreadMetrics>& threads, 
                                          const std::vector<int>& available_cores, 
                                          int numa_id);
    bool perform_activity_based_binding(std::vector<ThreadMetrics>& threads, 
                                       const std::vector<std::pair<int, int>>& physical_cores, 
                                       int numa_id);
    bool perform_miss_ratio_based_binding(std::vector<ThreadMetrics>& threads, 
                                         const std::vector<std::pair<int, int>>& physical_cores, 
                                         int numa_id);
    bool perform_hybrid_based_binding(std::vector<ThreadMetrics>& threads, 
                                     const std::vector<std::pair<int, int>>& physical_cores, 
                                     int numa_id);
    bool perform_complementary_binding(const std::vector<ThreadMetrics*>& group1, 
                                      const std::vector<ThreadMetrics*>& group2, 
                                      const std::vector<std::pair<int, int>>& physical_cores, 
                                      int numa_id, 
                                      const std::string& strategy_name);
    bool perform_simple_binding(const std::vector<ThreadMetrics>& threads, 
                               const std::vector<int>& available_cores, 
                               int numa_id);
    std::vector<std::pair<int, int>> get_physical_core_mapping(const std::vector<int>& available_cores);
    std::string get_thread_description(const ThreadMetrics& thread, const std::string& strategy_name);

private:
    // Core functionality
    bool initialize_system_info();
    bool collect_thread_metrics();
    bool perform_thread_binding();
    void binding_worker_thread();
    void collection_worker_thread();

    // Helper methods
    std::vector<int> get_thread_cpu_affinity(int tid);
    int get_numa_id_for_core(int core_id);
    std::vector<int> get_available_cores_for_numa(int numa_id);
    bool bind_thread_to_core(int tid, int core_id);
    std::vector<ThreadMetrics> find_threads_by_keyword(const std::string& keyword);
    void update_statistics(bool binding_success);
    void print_collection_summary(const std::vector<ThreadMetrics>& threads);
    void print_binding_analysis(const std::vector<ThreadMetrics>& threads, const std::string& analysis_type);
    bool perform_overflow_binding(const std::vector<const ThreadMetrics*>& remaining_threads,
                                 const std::vector<std::pair<int, int>>& physical_cores,
                                 size_t cores_used,
                                 int numa_id,
                                 const std::string& strategy_name,
                                 int& total_bindings,
                                 int& successful_bindings);

    // Member variables
    Config config_;
    std::atomic<bool> running_;
    std::atomic<bool> should_stop_;
    
    std::vector<ThreadMetrics> current_threads_;
    std::vector<NumaNode> numa_nodes_;
    std::vector<CpuCore> cpu_cores_;
    
    std::thread binding_thread_;
    std::thread collection_thread_;
    
    mutable std::mutex data_mutex_;
    mutable std::mutex stats_mutex_;
    std::condition_variable cv_;
    
    Statistics stats_;
    
    // PMU related
    int pmu_fd_;
    std::unique_ptr<PmuAttr> pmu_attr_;
};

} // namespace ThreadBinding

#endif // THREAD_BINDING_H 